import tkinter as tk
from tkinter import ttk
from ui.BasePage import BasePage

class ModelSettingPage(BasePage):
    """模型设置页面（整合版）"""

    def __init__(self, parent, controller):
        super().__init__(parent, controller)

        # 模型选择框架
        model_frame = ttk.LabelFrame(self.frame, text="模型选择")
        model_frame.pack(fill=tk.X, padx=10, pady=5)

        self.model_var = tk.StringVar(value="Informer-CBAM故障诊断模型")
        models = ["Informer-CBAM故障诊断模型", "Transformer-LSTM模型", "SDAE模型"]
        for model in models:
            rb = ttk.Radiobutton(
                model_frame, text=model, variable=self.model_var, value=model
            )
            rb.pack(anchor=tk.W, padx=10, pady=2)

        # 超参数设置框架
        hyper_param_frame = ttk.LabelFrame(self.frame, text="超参数设置")
        hyper_param_frame.pack(fill=tk.X, padx=10, pady=5)

        self.hyper_params = {}  # 存储超参数控件
        params = [
            ("批量大小：", 64, {"from_": 1, "to": 1000, "increment": 1, "format": "%d"}),
            ("学习率：", 0.001, {"from_": 0.0001, "to": 0.1, "increment": 0.0001, "format": "%.4f"}),
            ("训练轮数：", 50, {"from_": 1, "to": 1000, "increment": 1, "format": "%d"}),
            ("Dropout率：", 0.5, {"from_": 0, "to": 1, "increment": 0.05, "format": "%.2f"}),
        ]

        for param, default, spinbox_opts in params:
            row_frame = ttk.Frame(hyper_param_frame)
            row_frame.pack(fill=tk.X, padx=5, pady=2)

            # 标签
            ttk.Label(row_frame, text=param).pack(side=tk.LEFT, padx=2, pady=2)

            # 微调框（合并输入和调整功能）
            spinbox = ttk.Spinbox(
                row_frame,
                from_=spinbox_opts["from_"],
                to=spinbox_opts["to"],
                increment=spinbox_opts["increment"],
                format=spinbox_opts.get("format", ""),
                width=12
            )
            spinbox.pack(side=tk.LEFT, fill=tk.X, expand=True, padx=2)
            spinbox.insert(0, default)
            self.hyper_params[param] = spinbox

        # 优化算法选择框架
        optimizer_frame = ttk.LabelFrame(self.frame, text="优化算法选择")
        optimizer_frame.pack(fill=tk.X, padx=10, pady=5)

        self.optimizer_var = tk.StringVar(value="麻雀优化算法（SSA）")
        optimizers = [
            ("麻雀优化算法（SSA）", self.optimizer_var, "SSA"),
            ("鲸鱼优化算法", self.optimizer_var, "WOA"),
        ]

        for text, var, value in optimizers:
            rb = ttk.Radiobutton(optimizer_frame, text=text, variable=var, value=value)
            rb.pack(anchor=tk.W, padx=10, pady=2)

        # 损失函数选择框架
        loss_frame = ttk.LabelFrame(self.frame, text="损失函数选择")
        loss_frame.pack(fill=tk.X, padx=10, pady=5)

        self.loss_var = tk.StringVar(value="交叉熵")
        losses = [("交叉熵", self.loss_var, "CE"), ("Focal Loss", self.loss_var, "FL")]

        for text, var, value in losses:
            rb = ttk.Radiobutton(loss_frame, text=text, variable=var, value=value)
            rb.pack(anchor=tk.W, padx=10, pady=2)

        # 模型初始化按钮 - 使用BasePage中的标准按钮样式
        self.init_btn = ttk.Button(
            self.frame, text="模型初始化", style="Info.TButton"
        )
        self.init_btn.pack(fill=tk.X, padx=10, pady=5)
        self.init_btn.bind("<Button-1>", self.on_model_init)

        # 信息显示文本框
        self.info_text = tk.Text(self.frame, height=6, wrap=tk.WORD)
        self.info_text.pack(fill=tk.X, padx=10, pady=5)
        self.info_text.insert(
            tk.END,
            "模型：Informer - CBAM故障诊断模型\n输入维度：1024\n输出维度：4\n设备：cuda\n总参数量：1,057,430\n可训练量：1,057,430",
        )
        self.info_text.config(state=tk.DISABLED)

    def on_model_init(self, event):
        """模型初始化按钮点击事件"""
        model_name = self.model_var.get()
        params = self.get_hyperparameters()
        optimizer = self.optimizer_var.get()
        loss_fn = self.loss_var.get()

        # 更新信息显示
        self.info_text.config(state=tk.NORMAL)
        self.info_text.delete(1.0, tk.END)
        self.info_text.insert(
            tk.END,
            f"模型：{model_name}\n"
            f"批量大小：{params['批量大小：']}\n"
            f"学习率：{params['学习率：']}\n"
            f"训练轮数：{params['训练轮数：']}\n"
            f"Dropout率：{params['Dropout率：']}\n"
            f"优化算法：{optimizer}\n"
            f"损失函数：{loss_fn}\n"
            f"状态：模型已初始化",
        )
        self.info_text.config(state=tk.DISABLED)

    def get_hyperparameters(self):
        """获取超参数设置"""
        params = {}
        for label, spinbox in self.hyper_params.items():
            try:
                value = spinbox.get()
                # 根据格式转换类型
                if label in ["批量大小：", "训练轮数："]:
                    params[label] = int(value)
                else:
                    params[label] = float(value)
            except ValueError:
                # 设置默认值
                default = (
                    64
                    if label == "批量大小："
                    else (
                        0.001
                        if label == "学习率："
                        else 50 if label == "训练轮数：" else 0.5
                    )
                )
                params[label] = default
        return params

    def get_model_config(self):
        """获取完整模型配置"""
        return {
            "model_type": self.model_var.get(),
            "hyperparameters": self.get_hyperparameters(),
            "optimizer": self.optimizer_var.get(),
            "loss_function": self.loss_var.get(),
        }